(* Copyright © 2019 Ariadne Devos
   SPDX-License-Identifier: GPL-2.0 or GPL-3.0 *)

Require Import Coq.Setoids.Setoid.

(* Traces can be finite (like lists) or infintie (like streams).
   Equality of each element is 'eq', not an arbitrary
   equivalence relationship. *)
Section trace.

(* Presumably an event, can be anything *)
Variable A : Type.

CoInductive trace :=
| Done : trace
| More : A -> trace -> trace.

(* Witnesses of extensional equality *)
CoInductive trace_eq_proof : trace -> trace -> Prop :=
| EqDone : trace_eq_proof Done Done
| EqMore : forall a t u, trace_eq_proof t u -> trace_eq_proof (More a t) (More a u).

Axiom streameq_leibniz : forall t u, trace_eq_proof t u <-> t = u.

(* Trick Coq in reducing seq.

   Adapted from: adam.chlipala.net/cpdt/html/Coinductive.html. *)
Definition frob (t : trace) : trace :=
  match t with
  | Done => Done
  | More a t' => More a t'
  end.

Lemma frob_idem (t : trace) : t = frob t.
Proof.
  destruct t.
  all : reflexivity.
Qed.

Section Concatenation.

CoFixpoint seq (u : trace) (v : trace) :=
  match u with
  | Done => v
  | More s u' => More s (seq u' v)
  end.

(* Neutrality of nil *)

Lemma seq_nil_l (t : trace) : seq Done t = t.
Proof.
  set (t' := seq Done t).
  rewrite (frob_idem t).
  rewrite (frob_idem t').
  destruct t.
  all : (simpl; reflexivity).
Qed.

CoFixpoint seq_nil_r' (t : trace) : trace_eq_proof (seq t Done) t.
Proof.
  refine (match t with
    | Done => _
    | More a t' => _
    end).
  - set (u := seq Done Done).
    rewrite (frob_idem u).
    simpl. constructor.
  - set (u := seq (More a t') Done).
    rewrite (frob_idem u).
    simpl.
    constructor.
    + apply seq_nil_r'.
Qed.

Definition seq_nil_r (t : trace) : seq t Done = t.
Proof. rewrite <- streameq_leibniz. exact (seq_nil_r' t). Qed.

(* Associativity of seq *)
CoFixpoint seq_assoc' (t : trace) (u : trace) (v : trace) : trace_eq_proof (seq (seq t u) v) (seq t (seq u v)).
Proof.
  refine (match t with
    | Done => _
    | More a t' => _
  end).
  + repeat rewrite seq_nil_l.
    rewrite streameq_leibniz.
    reflexivity.
  + set (x := seq (More a t') u).
    set (y := seq (More a t') (seq u v)).
    rewrite (frob_idem x).
    rewrite (frob_idem y).
    simpl.
    clear x y.
    set (x := seq (More a (seq t' u)) v).
    rewrite (frob_idem x).
    simpl.
    constructor.
    apply seq_assoc'.
Qed.

Lemma seq_assoc (t : trace) (u : trace) (v : trace) : seq (seq t u) v = seq t (seq u v).
Proof. pose seq_assoc'. pose streameq_leibniz. firstorder. Qed.

(* Therefore, seq is a monoid. *)

(* induction lemma *)
Theorem seq_more_lift : forall a t u, seq (More a t) u = More a (seq t u).
Proof.
  intros.
  set (x := seq (More a t) u).
  rewrite (frob_idem x).
  cbn.
  reflexivity.
Qed.

End Concatenation.
End trace.

Section tmap.

CoFixpoint tmap A B (f : A -> B) (t : trace A) : trace B :=
  match t with
  | Done _ => Done _
  | More _ a t' => More _ (f a) (tmap A B f t')
  end.

Theorem tmap_empty A B f : tmap A B f (Done _) = Done _.
Proof.
  rewrite (frob_idem _ (tmap _ _ f (Done _))).
  cbn.
  reflexivity.
Qed.

Theorem tmap_peel A B (f : A -> B) (a : A) (t' : trace A) : tmap _ _ f (More _ a t') = More _ (f a) (tmap _ _ f t').
Proof.
  rewrite (frob_idem _ (tmap _ _ f (More _ _ _))).
  cbn.
  reflexivity.
Qed.

CoFixpoint tmap_compose' A B C (g : B -> C) (f : A -> B) (t : trace A) : trace_eq_proof _ (tmap _ _ g (tmap _ _ f t)) (tmap _ _ (Basics.compose g f) t).
Proof.
  destruct t.
  + repeat rewrite tmap_empty.
    constructor.
  + repeat rewrite tmap_peel.
    constructor.
    apply tmap_compose'.
Qed.

Theorem tmap_compose A B C (g : B -> C) (f : A -> B) (t : trace A) : tmap _ _ g (tmap _ _ f t) = tmap _ _ (Basics.compose g f) t.
Proof. rewrite <- streameq_leibniz. apply tmap_compose'. Qed.

CoFixpoint tmap_merge' A B (f : A -> B) t u : trace_eq_proof _ (seq _ (tmap _ _ f t) (tmap _ _ f u)) (tmap _ _ f (seq _ t u)).
Proof.
  destruct t.
  + rewrite tmap_empty.
    repeat rewrite seq_nil_l.
    rewrite streameq_leibniz.
    reflexivity.
  + rewrite seq_more_lift.
    repeat rewrite tmap_peel.
    rewrite seq_more_lift.
    constructor.
    apply tmap_merge'.
Qed.

Theorem tmap_merge A B (f : A -> B) t u : seq _ (tmap _ _ f t) (tmap _ _ f u) = tmap _ _ f (seq _ t u).
Proof.
  rewrite <- streameq_leibniz.
  apply tmap_merge'.
Qed.

End tmap.
